package ru.ifmo.ct.network.brafs.client;

import ru.ifmo.ct.network.brafs.IllegalCommandFormatException;
import ru.ifmo.ct.network.brafs.Utils;
import ru.ifmo.ct.network.brafs.file.SegmentId;
import ru.ifmo.ct.network.brafs.node.Command;
import ru.ifmo.ct.network.brafs.node.CommandType;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.*;
import java.nio.ByteBuffer;
import java.nio.charset.Charset;
import java.util.*;

/**
 * @author Alex Vikharev vikharev@yandex-team.ru
 *         created 20.04.12
 */
public class BraServerClient implements BraServer {

    private final String serverAddress;
    private final int serverPort;

    private Socket socket;
    private static final int MAX_SEGMENT_SIZE = 1024 * 1024;

    public BraServerClient(final String serverAddress, final int serverPort) {
        this.serverAddress = serverAddress;
        this.serverPort = serverPort;
    }

    @Override
    public void connect() throws IOException {
        socket = new Socket(serverAddress, serverPort);
    }

    @Override
    public void disconnect() throws IOException {
        socket.close();
    }

    private ServerSocket openSocketAtRandomPort() {
        final Random random = new Random();
        for (int i = 0; i < 10000; i++) {
            final int port = random.nextInt(65535 - 1024) + 1024;
            try {
                return new ServerSocket(port);
            } catch (IOException ignored) {
            }
        }
        throw new RuntimeException("No available ports");
    }

    private static class SegmentsServer implements Runnable {

        final ServerSocket server;
        final Map<SegmentId, byte[]> holder;
        volatile boolean allSegmentsSend;

        private SegmentsServer(final ServerSocket server) {
            this.server = server;
            holder = new HashMap<>();
        }

        void done() {
            allSegmentsSend = true;
        }

        @Override
        public void run() {
            try {
                server.setSoTimeout(500);
            } catch (SocketException e) {
                System.err.println("Error setting SO_TIMEOUT: " + e.getLocalizedMessage());
            }
            while (!Thread.interrupted() && !(allSegmentsSend && holder.isEmpty())) {
                try {
                    final Socket client = server.accept();
                    new Thread(new ClientProcessor(holder, client)).start();
                } catch (IOException e) {
                    if (!(e instanceof SocketTimeoutException)) {
                        System.err.println("Error: " + e.getLocalizedMessage());
                    }
                }
            }
        }
    }

    private static class ClientProcessor implements Runnable {

        private final Socket client;
        private final Map<SegmentId, byte[]> holder;

        private ClientProcessor(final Map<SegmentId, byte[]> holder, final Socket client) {
            this.holder = holder;
            this.client = client;
        }

        @Override
        public void run() {
            final Command command;
            try {
                command = Command.read(client.getInputStream());
            } catch (Exception e) {
                System.err.println("Error reading command from client: " + e.getLocalizedMessage());
                return;
            }

            final String fileName = command.getFileName();
            final long segmentId = command.getSegmentId();
            final SegmentId key = new SegmentId(fileName, segmentId);

            if (command.getType() == CommandType.GET_SEGMENT) {
                final byte[] bytes = holder.get(key);
                if (bytes != null) {
                    final ByteBuffer buffer = ByteBuffer.allocate(bytes.length + 4);
                    buffer.putInt(bytes.length);
                    buffer.put(bytes);
                    try {
                        final OutputStream outputStream = client.getOutputStream();
                        outputStream.write(buffer.array());
                        client.close();
                        holder.remove(key);
                    } catch (IOException e) {
                        System.err.println("Error write response: " + e.getLocalizedMessage());
                    }
                }
            } else {
                System.err.println("Unexpected command: " + command.getType());
            }
        }
    }

    @Override
    public OutputStream getOutputStream(final String fileId) {
        final ServerSocket server = openSocketAtRandomPort();
        final int segmentsServerPort = server.getLocalPort();
        final SegmentsServer segmentsServer = new SegmentsServer(server);
        final Thread segmentsServerThread = new Thread(segmentsServer);
        segmentsServerThread.start();

        return new OutputStream() {
            byte[] segment = new byte[MAX_SEGMENT_SIZE];
            int position = 0;
            int segmentId = 0;

            @Override
            public void write(final int b) throws IOException {
                if (segment.length == position) {
                    sendSegment();
                    segment = new byte[MAX_SEGMENT_SIZE];
                }
                segment[position++] = (byte) (b & 0xFF);
            }

            @Override
            public void close() throws IOException {
                if (position != 0) {
                    sendSegment();
                }
                segmentsServer.done();
                System.err.println("All segments written: " + segmentId);
            }

            private void sendSegment() throws IOException {
                final Command command = new Command(CommandType.ADD_SEGMENT, new byte[4], segmentsServerPort, segmentId, fileId);
                final byte[] actualSegment = new byte[position];
                System.arraycopy(segment, 0, actualSegment, 0, position);
                segmentsServer.holder.put(new SegmentId(fileId, segmentId), actualSegment);
                command.write(socket.getOutputStream());
                position = 0;
                segmentId++;
            }
        };
    }

    @Override
    public InputStream getInputStream(final String fileId) {
        return new InputStream() {
            byte[] segment = null;
            int position = 0;
            int segmentId = 0;
            boolean end = false;

            @Override
            public int read() throws IOException {
                if (end) {
                    return -1;
                }
                if (segment == null || segment.length == position) {
                    final Command command = new Command(CommandType.GET_SEGMENT, segmentId, fileId);
                    command.write(socket.getOutputStream());
                    try {
                        final Command responseCommand = Command.read(socket.getInputStream());
                        if (responseCommand.getType() == CommandType.GET_SEGMENT_FROM) {
                            segment = fetchSegment(command.getFileName(), segmentId, command.getIp(), command.getPort());
                            segmentId++;
                        } else {
                            end = true;
                            return -1;
                        }
                    } catch (IllegalCommandFormatException e) {
                        end = true;
                        return -1;
                    }
                }
                return segment[position++] & 0xFF;
            }
        };
    }


    @Override
    public List<String> listFiles(final String prefix) throws IOException {
        final Command command = new Command(CommandType.LIST_FILES, 0, prefix);
        command.write(socket.getOutputStream());
        final InputStream inputStream = socket.getInputStream();
        final ByteBuffer buffer = Utils.read(inputStream, 4);
        final int n = buffer.getInt(0);
        final List<String> result = new ArrayList<>(n);
        for (int i = 0; i < n; i++) {
            final int len = Utils.readInt(inputStream);
            final byte[] bytes = Utils.read(inputStream, len).array();
            result.add(new String(bytes, Charset.forName("utf-8")));
        }
        return result;
    }

    private static byte[] fetchSegment(final String fileName,
                                       final long segmentId,
                                       final byte[] ip,
                                       final int port) throws IOException {
        try (Socket socket = new Socket(InetAddress.getByAddress(ip), port)) {
            final Command getCommand = new Command(CommandType.GET_SEGMENT, segmentId, fileName);
            getCommand.write(socket.getOutputStream());
            final InputStream inputStream = socket.getInputStream();

            final ByteBuffer buffer = Utils.read(inputStream, MAX_SEGMENT_SIZE + 4);

            final int len = buffer.getInt(0);
            if (buffer.position() != len + 4) {
                throw new IOException("Unexpected segment length. Expected: " + (len + 4) + " Got: " + buffer.position());
            }
            final byte[] data = new byte[len];
            buffer.get(data, 4, len);
            return data;
        }

    }
}
